Skip to content

Enabling Geoinfo Feature for Anemoi Datasets #1634

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
wael-mika wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

Enabling Geoinfo Feature for Anemoi Datasets #1634

wael-mika wants to merge 3 commits into from
+385 −6

Conversation

@wael-mika
Copy link
Contributor

@wael-mika wael-mika commented Jan 16, 2026 •
edited by clessig
Loading

Description

Geoinfo channels are static/constant-in-time variables such as altitude, land-sea mask, orography, etc. These variables provide geographic context to the model and can be used in the decoder for location-aware predictions.

Previously, geoinfo was only supported for OBS-type streams. This implementation extends geoinfo support to anemoi-type streams (ERA5, CERRA, etc.).

brief summary of the changes introduced in this pull request.

  • Caching Logic: Added _cached_geoinfo to store data after the first read. Updated _get() to pull from this cache instead of hitting the disk repeatedly, and ensured init_empty() resets it.
  • Channel Selection: Added select_geoinfo_channels() to filter the config for static fields (is_constant_in_time=True).
  • Initialization: Updated init to wire everything together—populating the channel indices, calculating normalization stats (mean/stdev), and priming the cache.

Issue Number

Closes #1428

Checklist before asking for review

  • I have performed a self-review of my code
  • My changes comply with basic sanity checks:
    • I have fixed formatting issues with ./scripts/actions.sh lint
    • I have run unit tests with ./scripts/actions.sh unit-test
    • I have documented my code and I have updated the docstrings.
    • I have added unit tests, if relevant
  • I have tried my changes with data and code:
    • I have run the integration tests with ./scripts/actions.sh integration-test
    • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
    • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
  • I have informed and aligned with people impacted by my change:
    • for config changes: the MatterMost channels and/or a design doc
    • for changes of dependencies: the MatterMost software development channel

@clessig
Copy link
Collaborator

clessig commented Jan 16, 2026

Geoinfo channels do not have to be constant in time. They can be dynamic variables and the goal of the issue is that any channel can be a geoinfo channel. So no caching is possible.

@wael-mika
Copy link
Contributor Author

wael-mika commented Jan 16, 2026

Thanks for your comment @clessig. I will adapt accordingly.

clessig
clessig reviewed Jan 19, 2026
View reviewed changes
Copy link
Collaborator

@clessig clessig left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See comments.

src/weathergen/datasets/multi_stream_data_sampler.py

self.mini_epoch = 0

# Log geoinfo configuration per stream
Copy link
Collaborator

@clessig clessig Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please remove. This will be printed as part of the config

Copy link
Contributor Author

@wael-mika wael-mika Jan 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, I will

config/streams/era5_1deg_cerra/cerra.yml
@@ -0,0 +1,33 @@
# (C) Copyright 2024 WeatherGenerator contributors.
Copy link
Collaborator

@clessig clessig Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's please remove all the config files. We need to find a place for them that is not the main repo.

Copy link
Contributor Author

@wael-mika wael-mika Jan 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The configs are here only for the reviewer to test the work, I will remove them before submitting the PR

@wael-mika
Copy link
Contributor Author

wael-mika commented Jan 19, 2026

Geoinfo channels do not have to be constant in time. They can be dynamic variables and the goal of the issue is that any channel can be a geoinfo channel. So no caching is possible.

As some of geoinfo channels are constant in time, do you think we should check whether that and cache the channel if it is constant? or treat them all as non-constant in time?

@clessig
Copy link
Collaborator

clessig commented Jan 19, 2026 •
edited
Loading

Geoinfo channels do not have to be constant in time. They can be dynamic variables and the goal of the issue is that any channel can be a geoinfo channel. So no caching is possible.

As some of geoinfo channels are constant in time, do you think we should check whether that and cache the channel if it is constant? or treat them all as non-constant in time?

Let's treat them all as non-constant in time. The benefit would be so small that it's not worth the additional complexity.

@wael-mika
Copy link
Contributor Author

wael-mika commented Jan 20, 2026

How do you think we should handle the NaN values in the geoinfo channels @clessig?
For example in Cerra, the channels al, lsm, and orog contain a lot of NaN values which will effect the normalization which leads to "Skipping prediction" errors
My proposal is to mask nans with zeros in the base reader

@clessig
Copy link
Collaborator

clessig commented Jan 20, 2026

How do you think we should handle the NaN values in the geoinfo channels @clessig? For example in Cerra, the channels al, lsm, and orog contain a lot of NaN values which will effect the normalization which leads to "Skipping prediction" errors My proposal is to mask nans with zeros in the base reader

NaN values are handled later in the MultiStreamDataSampler. So nothing has to be done from your side.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@clessig clessig clessig left review comments

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

WeatherGen-dev
Status: No status

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@wael-mika @clessig